import os.path
from io import BytesIO

import cv2
import numpy as np
from PIL import ImageOps, Image, UnidentifiedImageError
from fastapi import APIRouter, HTTPException, UploadFile
from loguru import logger
from ocr_process.infer.predict import ClasPredict, TextImagePredict
from ocr_process.transformers import ImgParser
from ocr_process.utils.img_utils import np2base64, base64_pil
from pydantic import BaseModel, Field
from starlette.responses import HTMLResponse

from ocr_http.handler import inference_models
from ocr_http.ocr_exception import NoSupportException
from taskinfo.uv_pipeline import UvPipeline
from utils.config import settings
from threading import RLock

router = APIRouter()

lock = RLock()


class InferenceRequest(BaseModel):
    image_data: str = Field(alias="imageData")
    single_channel: bool = Field(default=False, alias="singleChannel")


class InferenceResponse(BaseModel):
    code: int = Field(default=0)
    message: str = Field(default="success")
    data: dict


clas_predict = ClasPredict(os.path.join(settings.path_config.clas_path, "inference.onnx"))
text_image_predict = TextImagePredict(os.path.join(settings.path_config.text_image_path, "inference.onnx"))


# uvdoc_predict = UvdocPredict(os.path.join(settings.path_config.uvdoc_path, "inference.onnx"))
uvdoc_predict = UvPipeline(os.path.join(settings.path_config.uvdoc_path, "inference.onnx"))


@router.get(path="/", summary="主页", response_class=HTMLResponse)
@router.get(path="/index", summary="主页", response_class=HTMLResponse)
@router.get(path="/index.html", summary="主页", response_class=HTMLResponse)
def root():
    with open("static/index.html", "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())


@router.post("/uploadfile")
def uploadfile(file: UploadFile):
    content = file.file.read()
    try:
        with BytesIO(content) as buffer:
            img = ImageOps.exif_transpose(Image.open(buffer).convert("RGB"))
            if img is None:
                raise HTTPException(status_code=400, detail="上传的图片无法识别，请确保上传的是有效图片文件")

    except (UnidentifiedImageError, OSError):
        raise HTTPException(status_code=400, detail="上传的图片无法识别，请确保上传的是有效图片文件")

    # ori_img = img.copy()
    image_np = np.array(img)

    logger.info("start clas_predict")
    # 图像分类预测
    con_flag, class_id = clas_predict(image_np)
    logger.info(f"con_flag = {con_flag}, class_id = {class_id}")
    if not con_flag:
        raise HTTPException(status_code=400,
                            detail="目前不支持该类型图片识别，请重新上传【荣程钢铁磅单, 天津港磅单】")

    if class_id == "jian_yuan":
        raise HTTPException(status_code=400,
                            detail="目前该模型不支持该类型图片识别，请重新上传【荣程钢铁, 天津港磅单磅单】")

    logger.info("start text_image_predict")
    # 方向预测
    image_np = text_image_predict(image_np)
    if class_id == "rong_cheng":
        result = uvdoc_predict(image_np)
        image_data = result['output_imgs'][0]
        # image_data = image_data.transpose(1, 2, 0)
        cv2.imwrite('result.jpg', image_data.copy())
        # if result['output_imgs'][0]:
        if result['output_imgs']:
            image_np = result['output_imgs'][0]


    # 单通道
    # img = Image.fromarray(image_np)
    # r_channel = img.split()[0]
    # # r_channel.save("red.jpg")
    # # print(type(r_channel))
    # image_np = np.array(r_channel)

    # img = Image.fromarray(image_np)
    # #转换成灰度图
    # gray_image = img.convert('L')
    # image_np = np.array(gray_image)

    # img = Image.fromarray(image_np)
    # img.save("orig.jpg")

    image_data = np2base64(image_np)
    # data = {"class_id": class_id, "image_data": image_data}

    if class_id not in inference_models.keys():
        raise HTTPException(status_code=404, detail="class_id not found")
    ie = inference_models[class_id]
    result = ie(image_data)
    logger.info(f"ocr_result = {result}")
    ocr_result = [{k: v[0]["text"]} for k, v in result[0].items()]

    logger.info(ocr_result)
    buffer = ImgParser.write_image_with_results(image=image_data, result=result[0], return_image=True)
    return {"class_id": class_id, "ocr_result": ocr_result, "img_data": np2base64(buffer)}


@router.post("/inference2")
def create_upload_file2(file: UploadFile):
    content = file.file.read()
    try:
        with BytesIO(content) as buffer:
            img = ImageOps.exif_transpose(Image.open(buffer).convert("RGB"))
            if img is None:
                raise HTTPException(status_code=400, detail="上传的图片无法识别，请确保上传的是有效图片文件")

    except (UnidentifiedImageError, OSError):
        raise HTTPException(status_code=400, detail="上传的图片无法识别，请确保上传的是有效图片文件")

    # ori_img = img.copy()
    image_np = np.array(img)

    logger.info("start clas_predict")
    # 图像分类预测
    con_flag, class_id = clas_predict(image_np)
    logger.info(f"con_flag = {con_flag}, class_id = {class_id}")
    if not con_flag:
        raise HTTPException(status_code=400,
                            detail="目前不支持该类型图片识别，请重新上传【荣程钢铁磅单, 天津港磅单】")

    if class_id == "jian_yuan":
        raise HTTPException(status_code=400,
                            detail="目前该模型不支持该类型图片识别，请重新上传【荣程钢铁, 天津港磅单磅单】")

    logger.info("start text_image_predict")
    # 方向预测
    image_np = text_image_predict(image_np)
    image_data = np2base64(image_np)
    # data = {"class_id": class_id, "image_data": image_data}

    if class_id not in inference_models.keys():
        raise HTTPException(status_code=404, detail="class_id not found")
    ie = inference_models[class_id]
    result = ie(image_data)
    logger.info(f"ocr_result = {result}")
    ocr_result = [{k: v[0]["text"]} for k, v in result[0].items()]

    logger.info(ocr_result)
    return {"class_id": class_id, "ocr_result": ocr_result}


def perform_text_image_predict(image_np: np.ndarray) -> np.ndarray:
    image_np = text_image_predict(image_np)
    logger.info("text_image_predict completed")
    return image_np


def perform_clas_predict(image_np: np.ndarray) -> int:
    logger.info("start clas_predict")
    con_flag, class_id = clas_predict(image_np)
    if not con_flag:
        # raise HTTPException(status_code=400, detail="目前不支持该类型图片识别，请重新上传【荣程钢铁磅单, 天津港磅单】")
        raise NoSupportException(message="目前不支持该类型图片识别，请重新上传【荣程钢铁磅单, 天津港磅单】")
    logger.info(f"clas_predict: con_flag = {con_flag}, class_id = {class_id}")
    return class_id


@router.post("/zy/inference")
def inference(request: InferenceRequest):
    lock.acquire()
    try:
        try:
            img = base64_pil(request.image_data)
        except:
            raise HTTPException(status_code=400, detail="上传的图片无法识别，请确保上传的是有效图片文件")
        image_np = np.array(img)
        # 图像分类预测
        class_id = perform_clas_predict(image_np)
        # 方向预测
        image_np = perform_text_image_predict(image_np)
        if class_id == "rong_cheng":
            result = uvdoc_predict(image_np)
            if result['output_imgs']:
                image_np = result['output_imgs'][0]
        if request.single_channel:
            logger.info("start single channel")
            img = Image.fromarray(image_np)
            r_channel = img.split()[0]
            image_np = np.array(r_channel)
        image_data = np2base64(image_np)
        if class_id not in inference_models.keys():
            # raise HTTPException(status_code=404, detail="class_id not found")
            raise NoSupportException(message="目前不支持该类型图片识别，请重新上传【荣程钢铁磅单, 天津港磅单】")
        ie = inference_models[class_id]
        result = ie(image_data)
        logger.info(f"ocr_result = {result}")
        ocr_result = [{k: v[0]["text"]} for k, v in result[0].items()]
        ocr_result = {key: value for item in ocr_result for key, value in item.items()}
        logger.info(ocr_result)
        return InferenceResponse(data={"class_id": class_id, "ocr_result": ocr_result})
    finally:
        lock.release()
